
from Scripts.env import Env
import config as cfg
from Scripts.keyboard import KeyboardCatch
from Scripts.joystick import Joystick
import os
import time
import copy

reset = False
end = False
car_vel_target = [0.0] * 2
car_vel_target_last = [0.0] * 2
car_vel_range = [
    [-cfg.BASE["max_linear_vel"], cfg.BASE["max_linear_vel"]],
    [-cfg.BASE["max_angular_vel"], cfg.BASE["max_angular_vel"]]
]
step = 0.2

def show(
    car_linear_vel_now: float, 
    car_angular_vel_now: float,
    wheel_vel_target: list,
    wheel_vel_now: list,
    joystick_exist: bool
):
    os.system("clear")
    print('*' * 20 + " 指令 " + "*" * 20)
    if joystick_exist:
        print("左摇杆前后 -> 前进/后退")
        print("右摇杆左右 -> 左转/右转")
        print("按钮A/B/X -> 停止/重置/关闭程序")
    else:
        print("W/S -> 前进/后退")
        print("Q/E -> 左转/右转")
        print("X/C/Esc -> 停止/重置/关闭程序")
    
    print("\n" + "-" * 20 + " car速度 " + "-" * 20)
    print("\t\t  linear_vel" + " " * 5 + "angular_vel")
    print(f"\ttarget: {car_vel_target[0]: 5.1f}" + " " * 10 + f"{car_vel_target[1]: 5.1f}")
    print(f"\tnow   : {car_linear_vel_now: 5.1f}" + " " * 10 + f"{car_angular_vel_now: 5.1f}")
    
    print("\n" + "-" * 20 + " 轮速 " + "-" * 20)
    print("\t" + " " * 10 +  "FL" + " " * 4 + "FR" + " " * 4 + "RL" + " " * 4 + "RR")
    print("\ttarget:", end="")
    for vel in wheel_vel_target:
        print(f"{vel: 5.1f} ", end="")
    
    print("\n\tnow:   ", end="")
    for vel in wheel_vel_now:
        print(f"{vel: 5.1f} ", end="")
    print()

def handle_keyboard(
    key_states: dict, 
):
    global reset, end, car_vel_target, car_vel_range
    
    if key_states["w"]:
        car_vel_target[0] += step
    elif key_states["s"]:
        car_vel_target[0] -= step
    elif key_states["q"]:
        car_vel_target[1] += step
    elif key_states["e"]:
        car_vel_target[1] -= step
    elif key_states["x"]:
        car_vel_target = [0, 0]
    elif key_states["c"]:
        car_vel_target = [0, 0]
        reset = True
    elif key_states["esc"]:
        end = True
    car_vel_target = limit_command(car_vel_target, car_vel_range)
    return car_vel_target
def limit_command(command, command_range):
    for i in range(len(command)):
        command[i] = min(max(command[i], command_range[i][0]), command_range[i][1])
    return command

def handle_joystick(
    axis_data: list,
    button_data: list,
):
    global reset, end, car_vel_target, car_vel_range, car_vel_target_last
    
    car_vel_target = [
        axis_data[1] * (car_vel_range[0][1] if axis_data[1] > 0 else abs(car_vel_range[0][0])),
        axis_data[2] * (car_vel_range[1][1] if axis_data[2] > 0 else abs(car_vel_range[1][0]))
    ]
    for i in range(1):
        if abs(car_vel_target[i] - car_vel_target_last[i]) > step:
            car_vel_target[i] = car_vel_target_last[i] + step if car_vel_target[i] > car_vel_target_last[i] else car_vel_target_last[i] - step
    car_vel_target_last = copy.deepcopy(car_vel_target)
    if button_data[0]:
        car_vel_target = [0, 0]
    elif button_data[1]:
        car_vel_target = [0, 0]
        reset = True
    elif button_data[3]:
        end = True
    
        
    car_vel_target = limit_command(car_vel_target, car_vel_range)
    return car_vel_target

if __name__ == "__main__":
    env = Env(
        xml_path = cfg.XML_PATH,
        enable_rangefinder_render = False
    )
    
    env.launch_viewer()
    
    joystick_exist = False
    joystick = Joystick()
    if joystick.init() is None:
        keyboardCatch = KeyboardCatch()
        keyboardCatch.run_keyboard_listener()
    else:
        joystick_exist = True
    
    while not end:
        if joystick_exist:
            car_vel_target = handle_joystick(
                joystick.get_axis_data(),
                joystick.get_button_data(),
            )
        else:
            car_vel_target = handle_keyboard(
                keyboardCatch.get_key_states(),
            )
        env.set_car_vel(car_vel_target)
        show(
            car_linear_vel_now=env.sensor.get_car_linear_vel(), 
            car_angular_vel_now=env.sensor.get_imu_sensor_data()[1][2],
            wheel_vel_target=env.controller.get_wheel_vel_target(),
            wheel_vel_now=env.sensor.get_wheel_vel_sensor_data(),
            joystick_exist=joystick_exist
        )
        # print(env.sensor.get_laser_data())
        time.sleep(0.1)
        

    env.close_viewer()
    